import sys
import os
import subprocess
from itertools import permutations
import parmed as pmd
import shutil
import numpy as np
import pandas as pd
from collections import defaultdict
from networkx.algorithms import isomorphism
from .paramgen import ParamGen
sys.path.append('..')
import coordmagic as cm
import simtk.openmm as mm
import simtk.openmm.app as app
from simtk.unit import kelvin, picoseconds, femtoseconds


class ParamRefine:
    '''update parameter for ParmED structure object from various sources '''
    def __init__(self,top_file = '',pmd_struct=''):
        self.refine_func = {'bonds':self.refine_bond_param,
                            'angles':self.refine_angle_param,
                            'coords':self.refine_eq,
                            }
        if top_file:
            self.load_file(top_file)
        if pmd_struct != '':
            self.convert2struct(pmd_struct)

    def load_file(self,top_file):
        filename = os.path.basename(top_file)
        self.basename,ext = os.path.splitext(filename)
        if ext in ['.log','.mol2','.pdb']:
            pg = ParamGen(filename)
            top,crd,chg = pg.gen_param()
            pmd_struct = pmd.load_file(top)
            pmd_coord = pmd.load_file(crd)
            pmd_struct.coordinates = pmd_coord.coordinates
        else:
            pmd_struct = pmd.load_file(top_file)
        self.convert2struct(pmd_struct)
        # self.ps = pmd_struct
        # self.ps = pmd.gromacs.GromacsTopologyFile.from_structure(pmd_struct)

    def convert2struct(self,st):
        self.ps = pmd.Structure()
        self.ps.atoms = st.atoms
        self.ps.residues = st.residues
        self.ps.bonds = st.bonds
        self.ps.angles = st.angles
        self.ps.dihedrals = st.dihedrals
        self.ps.bond_types = st.bond_types
        self.ps.angle_types = st.angle_types
        self.ps.dihedral_types = st.dihedral_types
        self.ps.combining_rule = st.combining_rule


    def gaulog2amber(self,logfile):
        '''convert gaussian log to amber promtop and inpcrd'''
        # first check if exe are in the path
        wkdir = self.basename + '_AMBER'
        try:
            os.mkdir(wkdir)
        except FileExistsError:
            pass
        cwd = os.getcwd()
        shutil.copy(logfile,wkdir)
        os.chdir(wkdir)
        com1 = 'antechamber -i {:s} -fi gout -o {:s}.mol2 -fo mol2 -c resp -at gaff2 -j 5 -pl 10 -dr no'.format(logfile,self.basename)
        try:
            subprocess.call(com1.split())
        except FileNotFoundError:
            print('Error!!! command antechamber not found. U should generate top file by yourself')
            sys.exit()
        com2 = 'parmchk2 -i {:s}.mol2 -f mol2 -o {:s}.frcmod -s gaff2 -a Y'.format(self.basename,self.basename)
        try:
            subprocess.call(com2.split())
        except FileNotFoundError:
            print('Error!!! command parmchk2 not found. U should generate top file by yourself')
            sys.exit()
        leapin = open(self.basename+'_leapin.in','w')
        leapin.write("loadamberparams {:s}.frcmod\n".format(self.basename))
        leapin.write("{:s}=loadmol2 {:s}.mol2\n".format(self.basename,self.basename))
        leapin.write("check {:s}\n".format(self.basename))
        leapin.write("saveamberparm {:s} {:s}.prmtop {:s}.inpcrd\n".format(self.basename,self.basename,self.basename))
        leapin.write("quit\n")
        leapin.close()
        com3 = "tleap -f {:s}_leapin.in".format(self.basename)
        try:
            subprocess.call(com3.split())
        except FileNotFoundError:
            print('Error!!! command tleap not found. U should generate top file by yourself')
            sys.exit()
        shutil.copy(self.basename+'.prmtop',cwd)
        shutil.copy(self.basename+'.inpcrd',cwd)
        os.chdir(cwd)


    def load_coord(self,coord_file):
        '''load coordinates from file and set to prm_structure.positions'''
        filename = os.path.basename(coord_file)
        basename,ext = os.path.splitext(filename)
        if ext in ['.gro','.pdb','.mol2','.pdbx','.inpcrd','.rst7','.restrt','.cif','.ncrst']:
            pmd_coord = pmd.load_file(coord_file)
            self.ps.coordinates = pmd_coord.coordinates
        elif ext in ['.res','.xyz','.log','.fchk','.xsf']:
            cm_st = cm.read_structure(coord_file)
            self.ps.coordinates = cm_st.coord
        else:
            print('format {:s} not suported'.format(ext))
            sys.exit()


    def refine_by(self,refine_object, refine_param = 'all'):
        if refine_param == 'all':
            refine_param = list(self.refine_func.keys())
        else:
            refine_param = refine_param.split(',')
        for k,v in refine_object.params.items():
            if k in refine_param:
                self.refine_func[k](v)

    def refine_charge(self,charge_file):
        chg_list = []
        cf = open(charge_file,'r')
        for i,c in enumerate(cf):
            atom = self.ps.atoms[i]
            old_chg =atom.charge
            new_chg = float(c.split()[-1])
            chg_list.append([atom.element,old_chg,new_chg])
            atom.charge = new_chg
        df = pd.DataFrame(chg_list,columns=['elem','old_chg','new_chg'])
        df['delta_chg'] = df['new_chg'] - df['old_chg']
        report = self.do_statistics(df,['delta_chg'],groupby=None,transpose=True)
        print(report)

    def adj_dihedral(self,contract='auto',param_struct=''):
        '''contract all dihedral types to dihedral type list '''
        if not param_struct:
            param_struct=self.ps
        d2dtype = defaultdict(str)
        for d in param_struct.dihedrals:
            a1 = d.atom1.idx
            a2 = d.atom2.idx
            a3 = d.atom3.idx
            a4 = d.atom4.idx
            d_idx = (a1,a2,a3,a4)
            atoms = (d.atom1,d.atom2,d.atom3,d.atom4)
            if type(d.type) is pmd.DihedralType:
                typelist = [d.type]
            elif type(d.type) is pmd.DihedralTypeList:
                typelist = [t for t in d.type]
            if  d_idx not in d2dtype:
                d2dtype[d_idx] =  {'typelist':typelist,'atoms':atoms,'ignore_end':d.ignore_end,'improper':d.improper}
            else:
                d2dtype[d_idx]['typelist'] += typelist
                d2dtype[d_idx]['ignore_end'] = d2dtype[d_idx]['ignore_end'] and d.ignore_end
        # remove d type with zero k and
        # for a,v in d2dtype.items():
        #     for dtype in v['typelist']:

        del param_struct.dihedrals[:]
        del param_struct.dihedral_types[:]
        param_struct.dihedral_types = pmd.TrackedList()
        param_struct.dihedrals = pmd.TrackedList()
        if contract == 'auto':
            for k,v in d2dtype.items():
                if len(v['typelist'])==1:
                    dtype = v['typelist'][0]
                    # dtype.list = param_struct.dihedral_types
                    inlist = 0
                    for d in param_struct.dihedral_types:
                        if type(d) is pmd.DihedralType:
                            if dtype == d:
                                dtype = d
                                inlist = 1
                    if inlist == 0:
                        param_struct.dihedral_types.append(dtype)
                    dihed = pmd.Dihedral(*v['atoms'],improper=v['improper'],ignore_end=v['ignore_end'],type=dtype)
                    param_struct.dihedrals.append(dihed)
                if len(v['typelist'])>1:
                    dtype = pmd.DihedralTypeList(v['typelist'], list=param_struct.dihedral_types)
                    inlist = 0
                    for d in param_struct.dihedral_types:
                        if type(d) is pmd.DihedralTypeList:
                            if dtype == d:
                                dtype = d
                                inlist = 1
                    if inlist == 0:
                        param_struct.dihedral_types.append(dtype)
                    dihed = pmd.Dihedral(*v['atoms'],improper=v['improper'],ignore_end=v['ignore_end'],type=dtype)
                    param_struct.dihedrals.append(dihed)

        if contract == 'all':
            for k,v in d2dtype.items():
                dtype = pmd.DihedralTypeList(v['typelist'],list=param_struct.dihedral_types)
                inlist = 0
                for d in param_struct.dihedral_types:
                    if dtype == d:
                        dtype = d
                        inlist =1
                if inlist == 0:
                    param_struct.dihedral_types.append(dtype)
                dihed = pmd.Dihedral(*v['atoms'], improper=v['improper'], ignore_end=v['ignore_end'], type=dtype)
                param_struct.dihedrals.append(dihed)

        if contract == 'none':
            for k,v in d2dtype.items():
                for dtype in v['typelist']:
                    dtype.list = param_struct.dihedral_types
                    inlist = 0
                    for d in param_struct.dihedral_types:
                        if dtype == d:
                            dtype = d
                            inlist = 1
                    if inlist == 0:
                        param_struct.dihedral_types.append(dtype)
                    dihed = pmd.Dihedral(*v['atoms'], improper=v['improper'], ignore_end=v['ignore_end'], type=dtype)
                    param_struct.dihedrals.append(dihed)

    def set_charge(self,target_charge = 0):
        charge = [a.charge for a in self.ps.atoms]
        total_q = sum(charge)
        delta_q = target_charge - total_q
        average_delta_q = delta_q / len(self.ps.atoms)
        print("The system have total charge of {:.6f}\n"
              "I will add charge {:.8f} to every atom to make system reach target charge {:.6f}"
              .format(total_q,average_delta_q,target_charge))
        for a in self.ps.atoms:
            a.charge += average_delta_q

    def compare_structure(self,coords1,coords2):
        '''compare coord with minimized stucture
        '''
        coords1 = np.array(coords1)/10
        bad1 = self.get_BAD(coords1)
        coords2 = np.array(coords2)/10
        bad2 = self.get_BAD(coords2)
        # bad_now,ene_now = self.get_minimized_bad()
        return self.compare_BAD(bad1,bad2)

    def transfer_param_from(self,res_obj,match='',atom='',param=''):
        match_sn = []
        atom_sn = []
        for i in match.split(','):
            if '-' in i:
                a = int(i.split('-')[0])
                b = int(i.split('-')[1]) + 1
                match_sn += list(range(a,b))
            else:
                match_sn.append(int(i))
        for i in atom.split(','):
            if '-' in i:
                a = int(i.split('-')[0])
                b = int(i.split('-')[1]) + 1
                atom_sn += list(range(a,b))
            else:
                atom_sn.append(int(i))

        G2 = res_obj['cm_frag']['graph'].subgraph(match_sn)

        st = cm.conver_structure(self.ps,'parmed')
        st.M.gen_molecules()
        nm = isomorphism.categorical_node_match("elem","")

        for id,m in st.molecules.items():
            print(m['formula'])
            G1 = m['graph']
            # print(len(G2.nodes.data()))
            GM = isomorphism.GraphMatcher(G1,G2,node_match=nm)
            mapping = []
            mapped_atom = []
            for m in list(GM.subgraph_isomorphisms_iter()):
               G1_atoms = ','.join(map(str,sorted([i for i in m.keys()])))
               if G1_atoms not in mapped_atom:
                   mapped_atom.append(G1_atoms)
                   mapping.append(m)
            print('{:d} match founded'.format(len(mapped_atom)))
        if len(mapped_atom) == 0:
            sys.exit()
        atom_map = []
        for m in mapping:
            atom_map.append({k:v for k,v in m.items() if v in atom_sn})
        self.adj_dihedral(contract='all')
        self.adj_dihedral(contract='all',param_struct=res_obj['pmd_struct'])
        mainB, mainA, mainD = self.get_BAD(get_param=True)
        auxB, auxA, auxD = self.get_BAD(param_struct=res_obj['pmd_struct'], get_param=True)
        for t in param.split(','):
            if t.startswith('b'):
                nm = 0
                for m in atom_map:
                    nm += 1
                    np = 0
                    for k,b in mainB.items():
                        if k[0] in m and k[1] in m:
                            if (m[k[0]],m[k[1]]) in auxB:
                                newb = auxB[(m[k[0]],m[k[1]])]
                            elif (m[k[1]],m[k[0]]) in auxB:
                                newb = auxB[(m[k[1]],m[k[0]])]
                            else:
                                print('Warning!!! param defined by {:s} in target do not have counterpart as {:s} in source'
                                        .format('-'.join([str(i) for i in k]), '-'.join([str(m[i]) for i in k])))
                                newb = None
                            if newb:
                                np += 1
                                b.type=newb.type
                print('{:d} X {:d} bond parameters transfered'.format(nm,np))
                self.ps.bond_types = pmd.TrackedList()
                for b in self.ps.bonds:
                    b.type.list = self.ps.bond_types
                    inlist = 0
                    for btype in self.ps.bond_types:
                        if b.type == btype:
                            b.type = btype
                            inlist = 1
                    if inlist == 0:
                        self.ps.bond_types.append(b.type)

            if t.startswith('a'):
                nm = 0
                for m in atom_map:
                    nm += 1
                    np = 0
                    for k,a in mainA.items():
                        if k[0] in m and k[1] in m and k[2] in m:
                            if (m[k[0]],m[k[1]],m[k[2]]) in auxA:
                                newa = auxA[(m[k[0]],m[k[1]],m[k[2]])]
                            elif (m[k[2]],m[k[1]],m[k[0]]) in auxA:
                                newa = auxA[(m[k[2]], m[k[1]], m[k[0]])]
                            else:
                                print('Warning!!! param defined by {:s} in target do not have counterpart as {:s} in source'
                                .format('-'.join([str(i) for i in k]), '-'.join([str(m[i]) for i in k])))
                                newa = None
                            if newa:
                                np += 1
                                a.type=newa.type
                print('{:d} X {:d} angle parameters transfered'.format(nm,np))
                self.ps.angle_types = pmd.TrackedList()
                for a in self.ps.angles:
                    a.type.list = self.ps.angle_types
                    inlist = 0
                    for atype in self.ps.angle_types:
                        if a.type == atype:
                            a.type = atype
                            inlist = 1
                    if inlist == 0:
                        self.ps.angle_types.append(a.type)

            if t.startswith('d'):
                nm = 0
                for m in atom_map:
                    nm += 1
                    np = 0
                    for k,d in mainD.items():
                        if k[0] in m and k[1] in m and k[2] in m and k[3] in m and not d.improper:
                            if (m[k[0]],m[k[1]],m[k[2]],m[k[3]]) in auxD:
                                newd = auxD[(m[k[0]],m[k[1]],m[k[2]],m[k[3]])]
                            elif (m[k[3]],m[k[2]],m[k[1]],m[k[0]]) in auxD:
                                newd = auxD[(m[k[3]],m[k[2]],m[k[1]],m[k[0]])]
                            else:
                                print('Warning!!! param defined by {:s} in target do not have counterpart as {:s} in source'
                                      .format('-'.join([str(i) for i in k]), '-'.join([str(m[i]) for i in k])))
                                newd = None
                            if newd:
                                np += 1
                                # print(k,d.type)
                                # print(k,newd.type)
                                d.type=newd.type
                print('{:d} X {:d} dihedral parameters transfered'.format(nm,np))
                self.ps.dihedral_types = pmd.TrackedList()
                for d in self.ps.dihedrals:
                    d.type.list = self.ps.dihedral_types
                    inlist = 0
                    for dtype in self.ps.dihedral_types:
                        if d.type == dtype:
                            d.type = dtype
                            inlist = 1
                    if inlist == 0:
                        self.ps.dihedral_types.append(d.type)

            if t.startswith('i'):
                nm = 0
                for m in atom_map:
                    nm += 1
                    np = 0
                    for k,d in mainD.items():
                        if k[0] in m and k[1] in m and k[2] in m and k[3] in m and d.improper:
                            newd = None
                            for c in permutations((m[k[0]],m[k[1]],m[k[2]],m[k[3]]),4):
                                if c in auxD:
                                    newd = auxD[c]
                                    np += 1
                                    d.type = newd.type
                            if not newd:
                                print('Warning!!! improper defined by {:s} in target do not have counterpart in source'
                                      .format('-'.join([str(i) for i in k])))
                print('{:d} X {:d} improper dihedral parameters transfered'.format(nm,np))
                self.ps.dihedral_types = pmd.TrackedList()
                for d in self.ps.dihedrals:
                    d.type.list = self.ps.dihedral_types
                    inlist = 0
                    for dtype in self.ps.dihedral_types:
                        if d.type == dtype:
                            d.type = dtype
                            inlist = 1
                    if inlist == 0:
                        self.ps.dihedral_types.append(d.type)

            if t.startswith('c'):
                for m in atom_map:
                    main_atom = m.keys()
                    aux_atom = m.values()
                    chg_main = [self.ps.atoms[i-1].charge for i in main_atom]
                    chg_aux = [res_obj['pmd_struct'].atoms[i-1].charge for i in aux_atom]
                    scale = sum(chg_main)/sum(chg_aux)
                    print('Origin charge: {:.3f}; replace charge: {:.3f}; scale factor: {:.3f}'
                          .format(sum(chg_main),sum(chg_aux),scale))
                    for k,v in m.items():
                        # print(self.ps.atoms[k-1].charge)
                        # print(res_obj['pmd_struct'].atoms[v-1].charge*scale)
                        self.ps.atoms[k-1].charge = res_obj['pmd_struct'].atoms[v-1].charge*scale
                print('{:d} X {:d} atom charge transfered'.format(len(atom_map),len(m)))


    def refine_bond_param(self,bond_param_dict):
        info_list = []
        new_bond_type = pmd.TrackedList()
        for bond in self.ps.bonds:
            r0 = bond.type.req
            k0 = bond.type.k
            a1 = str(bond.atom1.idx + 1)
            a2 = str(bond.atom2.idx + 1)
            if (a1, a2) in bond_param_dict:
                r, k = bond_param_dict[(a1, a2)]
                bond.type = pmd.BondType(k, r, new_bond_type)
            elif (a2, a1) in bond_param_dict:
                r, k = bond_param_dict[(a1, a2)]
                bond.type = pmd.BondType(k, r, new_bond_type)
            new_bond_type.append(bond.type)
            info_list.append(['-'.join([str(a1), str(a2)]), r0,r, k0,k])
        self.ps.bond_types = new_bond_type
        df = pd.DataFrame(info_list, columns=['atoms','old_eq','new_eq','old_k','new_k'])
        df['delta_eq'] = df['new_eq'] - df['old_eq']
        df['delta_k'] = df['new_k'] - df['old_k']
        print('Statistics on bond parameter refinement:')
        report = self.do_statistics(df,['delta_eq','delta_k'],groupby=None,transpose=True)
        print(report)
        try:
            os.mkdir('PARAM_REFINE') 
        except:
            pass
        dfname = os.path.join('PARAM_REFINE',self.basename+'_RBEK.csv')
        df.to_csv(dfname)
        print('Save the comparsion of old and new bond E,K to {:s}'.format(dfname))

    def do_statistics(self,df,col_name,groupby='type',transpose=False):
        def mae(x):
            return x.abs().mean()
        def rmse(x):
            return (x**2).mean()**0.5
        a=mae
        b=rmse
        agg_dict = {i:[mae,rmse,'mean','min','max','count'] for i in col_name}
        # print(agg_dict)
        if groupby:
            report = df.groupby(groupby).agg(agg_dict)
            report.columns = report.columns.droplevel()
        else:
            report = df.agg(agg_dict)
        if transpose:
            report = report.transpose()
        report.index.name = None
        return report

    def refine_angle_param(self,angle_param_dict):
        info_list = []
        new_angle_type = pmd.TrackedList()
        for angle in self.ps.angles:
            a0 = angle.type.theteq
            k0 = angle.type.k
            a1 = str(angle.atom1.idx + 1)
            a2 = str(angle.atom2.idx + 1)
            a3 = str(angle.atom3.idx + 1)
            if (a1, a2, a3) in angle_param_dict:
                a, k = angle_param_dict[(a1, a2, a3)]
                angle.type = pmd.AngleType(k, a, new_angle_type)
            elif (a3, a2, a1) in angle_param_dict:
                a, k = angle_param_dict[(a3, a2, a1)]
                angle.type = pmd.BondType(k, a, new_angle_type)
            new_angle_type.append(angle.type)
            info_list.append(['-'.join([str(a1), str(a2), str(a3)]), a0,a, k0,k])
        self.ps.angle_types = new_angle_type
        df = pd.DataFrame(info_list, columns=['atoms','old_eq','new_eq','old_k','new_k'])
        df['delta_eq'] = df['new_eq'] - df['old_eq']
        df['delta_k'] = df['new_k'] - df['old_k']
        print('Statistics on angle parameter refinement:')
        report = self.do_statistics(df,['delta_eq','delta_k'],groupby=None,transpose=True)
        print(report)
        try:
            os.mkdir('PARAM_REFINE') 
        except:
            pass
        dfname = os.path.join('PARAM_REFINE',self.basename+'_RAEK.csv')
        df.to_csv(dfname)
        print('Save the comparsion of old and new angle E,K to {:s}'.format(dfname))


    def refine_eq(self,coords):
        info_list = []
        bonds,angles,dihedrals = self.get_BAD(coords)
        new_bond_type= pmd.TrackedList()
        new_angle_type= pmd.TrackedList()
        for bond in self.ps.bonds:
            a1 = bond.atom1.idx + 1
            a2 = bond.atom2.idx + 1
            new_eq_len = bonds[(a1,a2)]
            current_eq_len = bond.type.req
            bond.type = pmd.BondType(bond.type.k, new_eq_len, new_bond_type)
            new_bond_type.append(bond.type)
            info_list.append(['bond','-'.join([str(a1),str(a2)]),current_eq_len,new_eq_len])
        for angle in self.ps.angles:
            a1 = angle.atom1.idx + 1
            a2 = angle.atom2.idx + 1
            a3 = angle.atom3.idx + 1
            new_eq_ang = angles[(a1,a2,a3)]
            current_eq_ang = angle.type.theteq
            angle.type = pmd.AngleType(angle.type.k, new_eq_ang, new_angle_type)
            new_angle_type.append(angle.type)
            info_list.append(['angle','-'.join([str(a1),str(a2),str(a3)]),current_eq_ang,new_eq_ang])
        self.ps.bond_types = new_bond_type
        self.ps.angle_types = new_angle_type
        df = pd.DataFrame(info_list, columns=['type', 'atoms','old_eq','new_eq'])
        df['delta_eq'] = df['new_eq'] - df['old_eq']
        print('Statistics on equilibrium bond length and angle refinement:')
        report = self.do_statistics(df,['delta_eq'])
        print(report)
        try:
            os.mkdir('PARAM_REFINE') 
        except:
            pass
        dfname = os.path.join('PARAM_REFINE',self.basename+'_RE.csv')
        df.to_csv(dfname)
        print('Save the comparsion of old and new equilibrium bond/angle to {:s}'.format(dfname))
        # pd.set_option('display.max_rows', None)

        # print(current_eq_len)
            # print(new_eq_len)
            # print(eq_len)
    #
    def minimize(self):
        system = self.ps.createSystem(nonbondedMethod=app.NoCutoff, constraints=app.HBonds)
        # Create the integrator to do Langevin dynamics
        integrator = mm.LangevinIntegrator(
            300 * kelvin,  # Temperature of heat bath
            1.0 / picoseconds,  # Friction coefficient
            2.0 * femtoseconds,  # Time step
        )

        sim = app.Simulation(self.ps.topology, system, integrator)
        # Set the particle positions
        sim.context.setPositions(self.ps.positions)
        sim.minimizeEnergy(tolerance=1)
        state = sim.context.getState(getEnergy=True)
        energy= state.getPotentialEnergy()
        positions = sim.context.getState(getPositions=True)
        coords = positions.getPositions(True)
        return coords,energy

    def get_minimized_bad(self):
        '''minimize structure and get internal coordinate in bad format'''
        ene,coords=self.minimize()
        bad = self.get_BAD(coords)
        return bad,ene

    def get_BAD(self,coords='', param_struct='', get_param=False):
        '''compute bond length angle and diheral for ps'''
        bonds=defaultdict(str)
        angles=defaultdict(str)
        dihedrals=defaultdict(str)
        bonds_param={}
        angles_param={}
        dihedrals_param={}
        if len(coords)==0:
            c = np.array(self.ps.coordinates)
        elif len(coords) == len(self.ps.coordinates):
            c = np.array(coords)
        else:
            print('The length of input coords {:d} do not match with coords in structures {:d}'
                  .format(len(coords),len(self.ps.coordinates)))
        if not param_struct:
            param_struct = self.ps
        for b in param_struct.bonds:
            coords1 = c[b.atom1.idx]
            coords2 = c[b.atom2.idx]
            dist = np.linalg.norm(coords1 - coords2)
            bonds[(b.atom1.idx+1,b.atom2.idx+1)] = dist
            bonds_param[(b.atom1.idx+1,b.atom2.idx+1)] = b
        for a in param_struct.angles:
            c1 = c[a.atom1.idx]
            c2 = c[a.atom2.idx]
            c3 = c[a.atom3.idx]
            v1 = c3 - c2
            v2 = c1 - c2
            cosang = np.dot(v1, v2)
            sinang = np.linalg.norm(np.cross(v1, v2))
            angle = np.rad2deg(np.arctan2(sinang, cosang))
            angles[(a.atom1.idx+1,a.atom2.idx+1,a.atom3.idx+1)] = angle
            angles_param[(a.atom1.idx + 1, a.atom2.idx + 1, a.atom3.idx + 1)] = a
        for d in param_struct.dihedrals:
            if d.improper == False:
                c1 = c[d.atom1.idx]
                c2 = c[d.atom2.idx]
                c3 = c[d.atom3.idx]
                c4 = c[d.atom4.idx]
                b0 = -1.0 *(c2-c1)
                b1 = c3 - c2
                b2 = c4 - c3
                b1 /= np.linalg.norm(b1)
                v = b0 - np.dot(b0, b1) * b1
                w = b2 - np.dot(b2, b1) * b1
                x = np.dot(v, w)
                y = np.dot(np.cross(b1, v), w)
                dihedral = np.degrees(np.arctan2(y, x))
                dihedrals[(d.atom1.idx+1,d.atom2.idx+1,d.atom3.idx+1,d.atom4.idx+1)] = dihedral
                dihedrals_param[(d.atom1.idx + 1, d.atom2.idx + 1, d.atom3.idx + 1, d.atom4.idx + 1)] = d
            else:
                dihedrals_param[(d.atom1.idx + 1, d.atom2.idx + 1, d.atom3.idx + 1, d.atom4.idx + 1)] = d
        # if param_struct is self.ps:
        #     for k,v in bonds_param.items():
        #         v.type.k=0
        #     for b in self.ps.bonds:
        #         print(b)

        if get_param:
            return bonds_param,angles_param,dihedrals_param
        else:
            return bonds,angles,dihedrals

    def compare_BAD(self,bad1,bad2):
        bond1,angle1,dihedral1 = bad1
        bond2,angle2,dihedral2 = bad2
        if list(bond1.keys()) == list(bond2.keys()):
            bond_diff = (np.array(list(bond1.values()))-np.array(list(bond2.values()))) * 10
            bond_info = zip(['-'.join(map(str,i)) for i in bond1.keys()],bond_diff)
            bond_info = [['bond'] + list(i) for i in bond_info]
        else:
            print('The bonds are different, so there will be no bond length comparison')
            bond_info = []
        if list(angle1.keys()) == list(angle2.keys()):
            angle_diff = np.array(list(angle1.values()))-np.array(list(angle2.values()))
            angle_info = list(zip(['-'.join(map(str,i)) for i in angle1.keys()],angle_diff))
            angle_info = [['angle'] + list(i) for i in angle_info]
        else:
            print('The angles are different, so there will be no angle comparison')
            angle_info = []
        if list(dihedral1.keys()) == list(dihedral2.keys()):
            raw_dihed_diff = np.array(list(dihedral1.values()))-np.array(list(dihedral2.values()))
            dihedral_diff = []
            for d in raw_dihed_diff:
                if d > 180:
                    d = 360-d
                if d < -180:
                    d = (d+360) * -1
                dihedral_diff.append(d)
            dihedral_diff = np.array(dihedral_diff)
            dihedral_info = list(zip(['-'.join(map(str,i)) for i in dihedral1.keys()],dihedral_diff))
            dihedral_info = [['dihedral'] + list(i) for i in dihedral_info]
        else:
            print('The dihedrals are different, so there will be no dihedral comparison')
            dihedral_info = []
        all_info = bond_info + angle_info + dihedral_info
        df = pd.DataFrame(all_info,columns=['type','atom','delta'])
        report = self.do_statistics(df,['delta'])
        return report

        # bond_mae = np.mean(np.abs(bond_diff))*10
        # bond_mse = np.mean(bond_diff)*10
        # bond_rmse = np.sqrt(np.mean(bond_diff**2))*10
        # bond_max = np.max(bond_diff)*10
        # bond_min = np.min(bond_diff)*10
        # angle_mae = np.mean(np.abs(angle_diff))
        # angle_mse = np.mean(angle_diff)
        # angle_rmse = np.sqrt(np.mean(angle_diff**2))
        # angle_max = np.max(angle_diff)
        # angle_min = np.min(angle_diff)
        # dihedral_mae = np.mean(np.abs(dihedral_diff))
        # dihedral_mse = np.mean(dihedral_diff)
        # dihedral_rmse = np.sqrt(np.mean(dihedral_diff**2))
        # dihedral_max = np.max(dihedral_diff)
        # dihedral_min = np.min(dihedral_diff)
        # info_list = [['bond',bond_mae,bond_rmse,bond_mse,bond_max,bond_min,len(bond_diff)]]
        # info_list.append(['angle',angle_mae,angle_rmse,angle_mse,angle_max,angle_min,len(angle_diff)])
        # info_list.append(['dihedral',dihedral_mae,dihedral_rmse,dihedral_mse,dihedral_max,dihedral_min,len(dihedral_diff)])
        # df = pd.DataFrame(info_list,columns=['type','mae','rmse','mse','max','min','count'])
        # df.set_index('type',inplace=True)
        # df.index.name = None
        # print(df)
